# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/29 17:20 
@Author  : ZhangShenao 
@File    : user_semantic_chunker.py 
@Desc    : 
"""

import dotenv
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_core.documents import Document
from langchain_experimental.text_splitter import SemanticChunker

dotenv.load_dotenv()

# create documents
page_content = """“But, soft! what light through yonder window breaks?
It is the east, and Juliet is the sun. 
Arise, fair sun, and kill the envious moon, Who is already sick and pale with grief, 
That thou her maid art far more fair than she: 
Be not her maid, since she is envious; 
Her vestal livery is but sick and green And none but fools do wear it; 
cast it off. For I am stony-broke, and true love doth cost: Money is alway ready; 
but who can value The dear purchased love?”
"""
docs = [Document(page_content=page_content, metadata={"source": "Romeo and Juliet.txt"})]

# create SemanticChunker
embeddings = ZhipuAIEmbeddings(model='embedding-2')
splitter = SemanticChunker(
    embeddings=embeddings,
    number_of_chunks=3,
    add_start_index=True,  # add start index
)

# split document to chunks
chunks = splitter.split_documents(docs)
for chunk in chunks:
    print(chunk.page_content)
    print(chunk.metadata)
